feat(ingest/kafka-connect): Kafka connect infer lineage from DataHub #15234
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
github-actions
bot
added
the
ingestion
❌ 1 Tests Failed:
View the top 1 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
treff7es
changed the title
| # Schema resolver configuration for enhanced lineage | ||
| use_schema_resolver: bool = Field( | ||
| default=False, | ||
| description="Use DataHub's schema metadata to enhance CDC connector lineage. " |
Contributor
There was a problem hiding this comment.
I would avoid introducing CDC as a new term to refer to Kafka Connect
CDC connector --> Kafka Connector
CDC sources/sinks -> Kafka Connect sources/sinks
datahub-cyborg
bot
added
pending-submitter-response
| It's used when connectors don't have table.include.list configured, meaning they | ||
| capture ALL tables from the database. | ||
| The method first tries to use cached URNs from SchemaResolver (populated from |
Contributor
There was a problem hiding this comment.
populated from previous ingestion runs? 🤔
| # Use graph.get_urns_by_filter() to get all datasets for this platform | ||
| # This is more efficient than a search query and uses the proper filtering API | ||
| all_urns = set( | ||
| self.schema_resolver.graph.get_urns_by_filter( |
Contributor
There was a problem hiding this comment.
We use the schema resolver here
all_urns = self.schema_resolver.get_urns()
`` `
* Is this Schema Resolver instance instantiated with the platform and env of the source database? anyway, either the resolver is populated with search such as the one below or it will be empty
and here
self.schema_resolver.graph.get_urns_by_filter(
entity_types=["dataset"],
platform=platform,
platform_instance=self.schema_resolver.platform_instance,
env=self.schema_resolver.env,
)
* here we just use the `graph` object, not actually using the Schema Resolver itself
I have the feeling that the usage of the Schema Resolver is very residual and only related to use its `DataHubGraph` object for doing the search.
Instead, I think it would be better to move down some of this pattern+search logic to the SchemaResolver itself. Or alternatively, just use the `DataHubGraph` object instead of making dependency with the SchemaResolver.
Comment on lines
+803
to
+805
| # Filter by platform | ||
| if f"dataPlatform:{platform}" not in urn: | ||
| continue |
Contributor
There was a problem hiding this comment.
this shouldn't be necessary with the get_urns_by_filter search, no?
Comment on lines
+807
to
+812
| # Filter by database - check if table_name starts with database prefix | ||
| if database_name: | ||
| if table_name.lower().startswith(f"{database_name.lower()}."): | ||
| # Remove database prefix to get "schema.table" | ||
| schema_table = table_name[len(database_name) + 1 :] | ||
| discovered_tables.append(schema_table) |
Contributor
There was a problem hiding this comment.
all this parsing in the URN.... I find it fragile
Couldn't we make a more specific search: eg search datasets in a given database (which should be a container)?
|
|
||
| # Build target URN using DatasetUrn helper with correct target platform | ||
| target_urn = DatasetUrn.create_from_ids( | ||
| platform_id=target_platform, |
|
|
||
| try: | ||
| # Get all URNs from schema resolver and filter for the source platform | ||
| # The cache may contain URNs from other platforms if shared across runs |
Contributor
There was a problem hiding this comment.
what do you mean "shared accross runs"?
Comment on lines
+1193
to
+1220
| if self.schema_resolver and self.schema_resolver.graph: | ||
| logger.info( | ||
| f"Kafka API unavailable for connector '{self.connector_manifest.name}' - " | ||
| f"querying DataHub for Kafka topics to expand pattern '{topics_regex}'" | ||
| ) | ||
| try: | ||
| # Query DataHub for all Kafka topics | ||
| kafka_topic_urns = list( | ||
| self.schema_resolver.graph.get_urns_by_filter( | ||
| platform="kafka", | ||
| env=self.schema_resolver.env, | ||
| entity_types=["dataset"], | ||
| ) | ||
| ) | ||
|
|
||
| datahub_topics = [] | ||
| for urn in kafka_topic_urns: | ||
| topic_name = self._extract_table_name_from_urn(urn) | ||
| if topic_name: | ||
| datahub_topics.append(topic_name) | ||
|
|
||
| matched_topics = matcher.filter_matches([topics_regex], datahub_topics) | ||
|
|
||
| logger.info( | ||
| f"Found {len(matched_topics)} Kafka topics in DataHub matching pattern '{topics_regex}' " | ||
| f"(out of {len(datahub_topics)} total Kafka topics)" | ||
| ) | ||
| return matched_topics |
Contributor
There was a problem hiding this comment.
We claim here using SchemaResolver but we are just using the grahp object and we are not even updating back SchemaResolver cache
As per my understanding, we should refresh the SchemaResolver cache with the search.
And ideally, we could move some logic down to the SchemaResolver class.
In SchemaResolver we have
def resolve_table(self, table: _TableName) -> Tuple[str, Optional[SchemaInfo]]: ...
maybe we could add some new method to resolve:
- urns and schemas for a given database
- or, urns and schemas for a given regexp pattern
I think pushing down this kind of functionality to SchemaResolver, that would simplify code in the source and would help to have better separation of responsibilities.
Comment on lines
+78
to
+98
| from datahub.sql_parsing.schema_resolver import SchemaResolver | ||
|
|
||
| # Get platform from connector instance (single source of truth) | ||
| platform = connector.get_platform() | ||
|
|
||
| # Get platform instance if configured | ||
| platform_instance = get_platform_instance( | ||
| config, connector.connector_manifest.name, platform | ||
| ) | ||
|
|
||
| logger.info( | ||
| f"Creating SchemaResolver for connector {connector.connector_manifest.name} " | ||
| f"with platform={platform}, platform_instance={platform_instance}" | ||
| ) | ||
|
|
||
| return SchemaResolver( | ||
| platform=platform, | ||
| platform_instance=platform_instance, | ||
| env=config.env, | ||
| graph=ctx.graph, | ||
| ) |
Contributor
There was a problem hiding this comment.
in bigquery, we do
return self.ctx.graph.initialize_schema_resolver_from_datahub(
platform=self.platform,
platform_instance=self.config.platform_instance,
env=self.config.env,
batch_size=self.config.schema_resolution_batch_size,
)
which instantiates SchemaResolver + populates its cache
that would make sense here too, right?
have you considered?
Contributor
There was a problem hiding this comment.
good abstractions here! ❤️
| ) | ||
|
|
||
|
|
||
| class TestParseCommaSeparatedList: |
Contributor
There was a problem hiding this comment.
nice tests covering all cases
wondering if parse_comma_separated_list may be moved to some utils class
| assert parse_comma_separated_list(input_str) == items | ||
|
|
||
|
|
||
| class TestConnectorConfigKeys: |
Contributor
There was a problem hiding this comment.
wondering the value of testing constants 😅
Contributor
sgomezvillamor
left a comment
sgomezvillamor
left a comment
There was a problem hiding this comment.
While I haven't reviewed all files in detail, overall looks pretty good
I haven't been able to check coverage in codecov, just may be high given the amount of tests
My concern is on the usage of SchemaResolver. We claim using it even in the user docs and this is reflected also in the configs, however the usage seems mostly about using the graph object in the SchemaResolver.
- we never refresh (or I missed it!) the internal cache of SchemaResolver (the caching would be one of the reasons to use SchemaResolver)
- and we still do a lot of "resolution" logic out of the SchemaResolver
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds automatic lineage inference from DataHub to the Kafka Connect source connector. Instead of relying solely on connector manifests, the ingestion can now query DataHub's metadata graph to resolve schemas and generate both table-level and column-level lineage.
Motivation
Currently, Kafka Connect lineage extraction is limited by what's explicitly declared in connector configurations. This PR enables:
table.include.list: "database.*"can now be resolved to actual table names by querying DataHubChanges
New Configuration Options
Added three new configuration fields to
KafkaConnectSourceConfig:Auto-Enable for Confluent Cloud
New behavior:
use_schema_resolveris automatically enabled when Confluent Cloud is detected via:confluent_cloud_environment_id+confluent_cloud_cluster_idconfigurationapi.confluent.cloud/connect/v1/)Users can opt-out by explicitly setting
use_schema_resolver: false.Core Components
SchemaResolver Integration (
connector_registry.py):create_schema_resolver()method to instantiate schema resolvers with platform-specific configurationsFine-Grained Lineage Extraction (
common.py):_extract_fine_grained_lineage()method inBaseConnectorFineGrainedLineageClassinstances for column-level lineageEnhanced Source Connectors (
source_connectors.py):ANALYTICS.PUBLIC.*→ actual tables)Pattern Matching (
pattern_matchers.py):database.*,schema.table*, etc.)Configuration Constants (
config_constants.py):Improved Topic Handling:
Code Quality Improvements
Usage
OSS Kafka Connect (Default Behavior)
Confluent Cloud (Auto-Enabled)
Confluent Cloud (Opt-Out)
OSS with Schema Resolver (Explicit Enable)
Testing
Test modules:
test_kafka_connect.py- Core connector teststest_kafka_connect_config_validation.py- Auto-enable logic tests (8 new tests)test_kafka_connect_schema_resolver.py- Schema resolver integrationtest_kafka_connect_snowflake_source.py- Snowflake connector teststest_kafka_connect_pattern_matchers.py- Pattern matching teststest_kafka_connect_config_constants.py- Configuration validationtest_kafka_connect_connector_registry.py- Connector registration testsDocumentation
docs/sources/kafka-connect/kafka-connect.mdwith:Breaking Changes
None - All features are opt-in (or auto-enabled only for Confluent Cloud). Existing Kafka Connect ingestions continue to work unchanged.
Default behavior:
use_schema_resolver: false(unchanged behavior)use_schema_resolver: true(auto-enabled, can be disabled)Prerequisites for Schema Resolver
IMPORTANT: For schema resolver to work, source database tables must be ingested into DataHub before running Kafka Connect ingestion. Without prior database ingestion, schema resolver will not find table metadata.
Recommended ingestion order:
🤖 Generated with Claude Code
Co-Authored-By: Claude [email protected]